package com.ruoyi.ade.service.impl;

import com.ruoyi.ade.domain.AdeKnowledge;
import com.ruoyi.ade.domain.AdeKnowledgeBase;
import com.ruoyi.ade.domain.AdeKnowledgeRecord;
import com.ruoyi.ade.domain.AdeLab;
import com.ruoyi.ade.llm.chunk.ChunkByMarkRequest;
import com.ruoyi.ade.llm.chunk.ChunkByNoneRequest;
import com.ruoyi.ade.llm.chunk.ChunkByPatternRequest;
import com.ruoyi.ade.llm.chunk.IAbstractChunkRequest;
import com.ruoyi.ade.llm.model.AdeRetrivedResponse;
import com.ruoyi.ade.llm.service.IAdeKnowledgeEmbedService;
import com.ruoyi.ade.mapper.AdeKnowledgeBaseMapper;
import com.ruoyi.ade.mapper.AdeKnowledgeMapper;
import com.ruoyi.ade.mapper.AdeKnowledgeRecordMapper;
import com.ruoyi.ade.model.AdeSearchRequest;
import com.ruoyi.ade.service.IAdeKnowledgeIndexService;
import com.ruoyi.common.utils.StringUtils;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;

@Service
@Slf4j
public class AdeKnowledgeIndexServiceImpl implements IAdeKnowledgeIndexService {

    @Autowired
    private AdeKnowledgeBaseMapper adeKnowledgeBaseMapper;  //ADE 向量库操作

    @Autowired
    private AdeKnowledgeMapper adeKnowledgeMapper;     //   ADE知识库 总库操作

    @Autowired
    private AdeKnowledgeRecordMapper adeKnowledgeRecordMapper;  //ADE 向量库中记录操作

    @Autowired
    private IAdeKnowledgeEmbedService adeKnowledgeEmbedService;   //ADE知识嵌入相关服务    包括索引和检索

    /**
     * 从ADE向量知识库中检索  ADE案例知识
     * @param request
     * @return
     */
    public AdeRetrivedResponse search(AdeSearchRequest request){
        AdeKnowledgeBase adeKnowledgeBase = adeKnowledgeBaseMapper.selectAdeKnowledgeBaseById(request.getAdekb());
        if (adeKnowledgeBase == null) throw new RuntimeException("No knowledge base found for:" + request.getAdekb());
        AdeRetrivedResponse response = adeKnowledgeEmbedService.search(adeKnowledgeBase.getCollectionName(), request.getEmbedModel(), request.getEmr(),request.getChunkType(), request.getSimScore(), request.getTopN(), false, request.isDebug());
        return response;
    }
    @Override
    public void execute(Long id) {
        //信息校验
        AdeKnowledgeBase adeKnowledgeBase = adeKnowledgeBaseMapper.selectAdeKnowledgeBaseById(id);
        if (adeKnowledgeBase == null) throw new RuntimeException("No knowledge base found for:" + id);

        if (!AdeLab.STATUS_READY.equals(adeKnowledgeBase.getStatus()) && !AdeLab.STATUS_FAILED.equals(adeKnowledgeBase.getStatus()))
            throw new RuntimeException("Current lab status exception:" + adeKnowledgeBase.getStatus());

        if (adeKnowledgeBase.getPosRate() < 0 || adeKnowledgeBase.getPosRate() > 1)
            throw new RuntimeException("Pos rate error:" + adeKnowledgeBase.getPosRate());
        if (adeKnowledgeBase.getSize() < 0) throw new RuntimeException("Size error:" + adeKnowledgeBase.getSize());

        log.info("Starting to index knowledge base:" + adeKnowledgeBase.getName());

        //更新状态 任务运行中
        updateStatus(id, AdeLab.STATUS_RUNNING, null, 0L, 0L, "");
        //启动任务
        new Thread(new Runnable() {

            @Override
            public void run() {
                //获取可用的ADE数据集
                List<AdeKnowledge> avis = adeKnowledgeMapper.selectAdeKnowledgeList(new AdeKnowledge());
                int avi_neg = 0;
                for (AdeKnowledge adeKnowledge : avis) {
                    if ("{}".equals(adeKnowledge.getAde())) avi_neg++;
                }
                int avi_pos = avis.size() - avi_neg;
                log.info("Lab id:{}, Pos avi size:{}, Neg avi size:{}, Total avi size:{}", id, avi_pos, avi_neg, avis.size());

                //确认需要索引的数据集
                long pos = (long) (adeKnowledgeBase.getPosRate() * adeKnowledgeBase.getSize());
                long neg = (adeKnowledgeBase.getSize() - pos);
                log.info("Lab id:{}, Pos size:{}, Neg size:{}, Total size:{}", id, pos, neg, adeKnowledgeBase.getSize());

                //对比分析
                if (adeKnowledgeBase.getSize() > avis.size()) {
                    log.warn("Lab id:{}, Required size is {}, But aviliable size is :{}", id, adeKnowledgeBase.getSize(), avis.size());
                    updateStatus(id, AdeLab.STATUS_FAILED, null, pos, neg, "Required size is " + adeKnowledgeBase.getSize() + ", But aviliable size is " + avis.size());
                    return;
                }
                if (pos > avi_pos) {
                    log.warn("Lab id:{}, Required Pos size is {}, But aviliable pos size is :{}", id, pos, avi_pos);
                    updateStatus(id, AdeLab.STATUS_FAILED, null, pos, neg, "Required Pos size is " + pos + ", But aviliable pos size is " + avi_pos);
                    return;  //结束
                }
                if (neg > avi_neg) {
                    log.warn("Lab id:{}, Required Neg size is {}, But aviliable neg size is :{}", id, neg, avi_neg);
                    updateStatus(id, AdeLab.STATUS_FAILED, null, pos, neg, "Required Neg size is " + neg + ", But aviliable neg size is " + avi_neg);
                    return;  //结束
                }

                //条件都满足了 则创建Mivlus Collection
                adeKnowledgeBase.setCollectionName("ADEKB_" + adeKnowledgeBase.getId() + "_" + adeKnowledgeBase.getSize());  //自动生成collection Name    ADEKB_3023_23
                try {
                    adeKnowledgeEmbedService.createCollection(adeKnowledgeBase.getCollectionName());
                    updateStatus(id, AdeLab.STATUS_RUNNING, adeKnowledgeBase.getCollectionName(), pos, neg, "Milvus create collection success");
                } catch (Exception e) {
                    updateStatus(id, AdeLab.STATUS_FAILED, adeKnowledgeBase.getCollectionName(),pos, neg, "Milvus create collection failed:" + e.getMessage());
                    return;
                }

                //条件合适了 可以分配pos neg了
                int posFinishCount = 0;
                int negFinishCount = 0;
                Collections.shuffle(avis); // 打乱顺序
                AdeKnowledge tmp = null;
                String collection = adeKnowledgeBase.getCollectionName();
                for (AdeKnowledge adeKnowledge : avis) {
                    tmp = null;
                    if ("{}".equals(adeKnowledge.getAde())) {  //如果是 neg
                        if (negFinishCount < neg) {
                            tmp = adeKnowledge;
                            negFinishCount++;
                        }
                    } else {                                    //如果是pos
                        if (posFinishCount < pos) {
                            tmp = adeKnowledge;
                            posFinishCount++;
                        }
                    }
                    if (tmp != null) {

                        //构建 ADE知识库的记录
                        AdeKnowledgeRecord record = new AdeKnowledgeRecord();
                        record.setBaseId(id);
                        record.setAdeKnowledgeId(tmp.getId());
                        record.setContent(tmp.getText());
                        record.setAdes(tmp.getAde());
                        record.setExplains(tmp.getExplain());
                        record.setStatus(AdeKnowledgeRecord.STATUS_OK);

                        try {
                            //构建索引记录内容  写入索引库
                            String explain = StringUtils.isEmpty(tmp.getExplain())?"" : tmp.getExplain();
                            com.ruoyi.ade.llm.model.AdeKnowledge adeKnowledge1 = new com.ruoyi.ade.llm.model.AdeKnowledge(String.valueOf(tmp.getId()), tmp.getText(), tmp.getAde(), explain);
                            //需要整理写入到索引库
                            adeKnowledgeEmbedService.index(collection, adeKnowledgeBase.getEmbedModel(), adeKnowledge1);
                        } catch (Exception e) {
                            //写入索引库失败
                            log.error("Ade Knowledge index:{} for AdeKnowledge ID:{},  Failed:{}", collection, tmp.getId(), e.getMessage());
                            record.setStatus(AdeKnowledgeRecord.STATUS_ERR);
                            record.setReason(e.getMessage());
                        }

                        //同步写入数据库
                        adeKnowledgeRecordMapper.insertAdeKnowledgeRecord(record);

                        //同步更新状态   更新进度
                        AdeKnowledgeBase update = new AdeKnowledgeBase();
                        update.setId(id);
                        update.setPosFinishSize((long) posFinishCount);
                        update.setNegFinishSize((long) negFinishCount);
                        adeKnowledgeBaseMapper.updateAdeKnowledgeBase(update);
                    }

                    if (posFinishCount >= pos && negFinishCount >= neg) {
                        updateStatus(id, AdeKnowledgeBase.STATUS_SUCCESS, null, null, null, null);
                        log.info("ADE KnowledgeBase id:{}, Finished update status:success", id);
                        break;  //如果都获取完了 则结束
                    }
                }

            }
        }).start();
    }

    private void updateStatus(Long id, String status, String collection, Long posSize, Long negSize, String message) {
        AdeKnowledgeBase adeKnowledgeBase = new AdeKnowledgeBase();
        adeKnowledgeBase.setId(id);
        if (collection != null) adeKnowledgeBase.setCollectionName(collection);
        if (posSize != null) adeKnowledgeBase.setPosSize(posSize);
        if (negSize != null) adeKnowledgeBase.setNegSize(negSize);
        if (message != null) adeKnowledgeBase.setReason(message);
        if (status != null) adeKnowledgeBase.setStatus(status);
        adeKnowledgeBaseMapper.updateAdeKnowledgeBase(adeKnowledgeBase);
    }


    public static void main(String[] args) {
        List<String> list = new ArrayList<>();
        list.add("a");
        list.add("b");
        list.add("C");
        list.add("d");
        list.add("e");
        Collections.shuffle(list);
        System.out.println(list);

        Collections.shuffle(list);
        System.out.println(list);

        Collections.shuffle(list);
        System.out.println(list);
    }
}
